import cv2
import numpy as np
import glob
import matplotlib.pyplot as plt
import os

# 设置棋盘格参数
square_size = 15  # 棋盘格每个方格的大小（单位：mm）
pattern_size = (11, 8)  # 棋盘格的行列数减一
images = glob.glob(r"C:\Users\Administrator\PycharmProjects\astra\captures\*.jpg")  # 修改为你的图像路径,*.  后面加图片型号

obj_points = []
img_points = []
objp = np.zeros((np.prod(pattern_size), 3), dtype=np.float32)
objp[:, :2] = np.mgrid[0:pattern_size[0], 0:pattern_size[1]].T.reshape(-1, 2) * square_size

det_success_num = 0

# scale = 0.5

# 打开文件以写入输出结果
with open("calibration_results.txt", "w") as file:
    # 遍历图像，检测棋盘格角点
    for i, image in enumerate(images):
        img = cv2.imread(image)
        if img is None:
            print(f"无法读取图像: {image}")
            file.write(f"无法读取图像: {image}\n")
            continue

        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        ret, corners = cv2.findChessboardCornersSB(gray, pattern_size)
        if ret:
            det_success_num += 1
            obj_points.append(objp)
            img_points.append(corners)
            cv2.drawChessboardCorners(img, pattern_size, corners, ret)
            cv2.imshow('img', img)
            cv2.waitKey(500)

            #-----缩放观察方向--------

            # display_img = cv2.resize(img, (0, 0), fx=scale, fy=scale)
            # cv2.imshow('Chessboard Detection (Scaled)', display_img)
            # key = cv2.waitKey(500)

    cv2.destroyAllWindows()

    # 相机标定
    ret, K, dist_coeffs, rvecs, tvecs = cv2.calibrateCamera(obj_points, img_points, gray.shape[::-1], None, None)
    file.write("相机内参矩阵 (K):\n")
    file.write(f"{K}\n")
    file.write("畸变系数:\n")
    file.write(f"{dist_coeffs}\n")

    # 计算并绘制重投影误差
    mean_error = 0
    errors = []
    error_x = []
    error_y = []

    for i in range(len(obj_points)):
        img_points_reproj, _ = cv2.projectPoints(obj_points[i], rvecs[i], tvecs[i], K, dist_coeffs)
        img_points_reproj = img_points_reproj.reshape(-1, 2)
        img_points_actual = img_points[i].reshape(-1, 2)

        error = np.linalg.norm(img_points_actual - img_points_reproj, axis=1)
        error_x.extend(img_points_actual[:, 0] - img_points_reproj[:, 0])
        error_y.extend(img_points_actual[:, 1] - img_points_reproj[:, 1])

        mean_error += np.mean(error)
        errors.append(np.mean(error))
        file.write(f"图像 {i + 1} 的重投影误差: {np.mean(error)}\n")

        # 标记重投影点
        img = cv2.imread(images[i])
        for j in range(len(img_points_actual)):
            pt_actual = tuple(img_points_actual[j].astype(int))
            pt_reproj = tuple(img_points_reproj[j].astype(int))
            cv2.circle(img, pt_actual, 5, (0, 255, 0), -1)  # 绿色：实际角点
            cv2.circle(img, pt_reproj, 5, (0, 0, 255), -1)  # 红色：重投影点
            cv2.line(img, pt_actual, pt_reproj, (255, 0, 0), 1)  # 蓝色连线

            # 添加图例
            legend = np.zeros((100, 200, 3), dtype=np.uint8)
            cv2.circle(legend, (30, 30), 5, (0, 255, 0), -1)
            cv2.putText(legend, 'Actual Points', (50, 35), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
            cv2.circle(legend, (30, 60), 5, (0, 0, 255), -1)
            cv2.putText(legend, 'Reprojected Points', (50, 65), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
            cv2.line(legend, (20, 80), (40, 80), (255, 0, 0), 1)
            cv2.putText(legend, 'Error Line', (50, 85), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

            img[10:110, -210:-10] = legend  # 将图例放在右上角

        cv2.imshow(f'Reprojection Error - Image {i + 1}', img)
        cv2.waitKey(1000)

    mean_error /= len(obj_points)
    file.write(f"总体平均重投影误差: {mean_error}\n")

    # 绘制重投影误差散点图
    plt.figure(figsize=(10, 6))
    plt.scatter(error_x, error_y, c='r', s=10, label='Reprojection Errors')
    plt.axhline(y=0, color='k', linestyle='--')
    plt.axvline(x=0, color='k', linestyle='--')
    plt.title('Reprojection Error Distribution')
    plt.xlabel('X Error (pixels)')
    plt.ylabel('Y Error (pixels)')
    plt.legend()
    plt.grid()
    plt.show()

    # 生成柱状图
    plt.figure(figsize=(10, 6))
    plt.bar(range(1, len(errors) + 1), errors, color='b')
    plt.title('Reprojection Errors for Each Image')
    plt.xlabel('Image Index')
    plt.ylabel('Reprojection Error (pixels)')
    plt.grid(True, axis='y')
    plt.show()

cv2.destroyAllWindows()

# 图片校正保存部分
# 创建保存校正后图像的目录
output_dir = "output/undistorted"
os.makedirs(output_dir, exist_ok=True)

for i, image in enumerate(images):
    img = cv2.imread(image)
    h, w = img.shape[:2]

    new_K, _ = cv2.getOptimalNewCameraMatrix(K, dist_coeffs, (w, h), 1, (w, h))
    undistorted_img = cv2.undistort(img, K, dist_coeffs, None, new_K)

    # 显示图像
    cv2.imshow('Original Image', img)
    cv2.imshow('Undistorted Image', undistorted_img)
    cv2.waitKey(1000)

    # 保存校正后的图像
    filename = os.path.basename(image)
    save_path = os.path.join(output_dir, f"undistorted_{filename}")
    cv2.imwrite(save_path, undistorted_img)

cv2.destroyAllWindows()

print("Camera Matrix K:\n", K)
print("Distortion Coefficients:\n", dist_coeffs)
